package com.greensam.pixelengine.websocket.handle;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.module.SimpleModule;
import com.fasterxml.jackson.databind.ser.std.ToStringSerializer;
import com.greensam.pixelengine.websocket.dto.WebSocketPushMessage;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 图片编辑 WebSocket 广播器
 * 负责会话管理和消息广播
 *
 * @author Ma Chengrui
 * @since 2025/11/4 16:46
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class PictureEditBroadcaster {

    // 保存所有连接的会话, key:pictureId value:session 集合
    private final Map<Long, Set<WebSocketSession>> PICTURE_SESSIONS = new ConcurrentHashMap<>();

    // 每张图片的编辑状态, key:pictureId value:当前正在编辑的用户ID
    private final Map<Long, Long> PICTURE_EDITING_USERS = new ConcurrentHashMap<>();

    private final ObjectMapper objectMapper;

    public PictureEditBroadcaster() {
        // 配置 ObjectMapper，使 Long 序列化为 String
        objectMapper = new ObjectMapper();
        SimpleModule module = new SimpleModule();
        module.addSerializer(Long.class, ToStringSerializer.instance);
        module.addSerializer(Long.TYPE, ToStringSerializer.instance);
        this.objectMapper.registerModule(module);
    }

    public void addSession(Long pictureId, WebSocketSession session) {
        PICTURE_SESSIONS.putIfAbsent(pictureId, ConcurrentHashMap.newKeySet());
        PICTURE_SESSIONS.get(pictureId).add(session);
    }

    public void removeSession(Long pictureId, WebSocketSession session) {
        Set<WebSocketSession> sessions = PICTURE_SESSIONS.get(pictureId);
        if (sessions != null) {
            sessions.remove(session);
            if (sessions.isEmpty()) {
                PICTURE_SESSIONS.remove(pictureId);
            }
        }
    }

    public void setEditingUser(Long pictureId, Long userId) {
        if (userId == null) {
            PICTURE_EDITING_USERS.remove(pictureId);
        } else {
            PICTURE_EDITING_USERS.put(pictureId, userId);
        }
    }

    public Long getEditingUser(Long pictureId) {
        return PICTURE_EDITING_USERS.get(pictureId);
    }

    public void broadcast(Long pictureId, WebSocketPushMessage<Map<String, Object>> message) {
        broadcast(pictureId, message, null);
    }

    /**
     * originSessionId 用于排除发起者（如编辑动作不需要发回 origin）
     */
    public void broadcast(Long pictureId, WebSocketPushMessage<Map<String, Object>> message, String originSessionId) {
        Set<WebSocketSession> sessions = PICTURE_SESSIONS.get(pictureId);
        if (CollectionUtils.isEmpty(sessions)) {
            return;
        }
        try {
            String json = objectMapper.writeValueAsString(message);
            TextMessage text = new TextMessage(json);
            for (WebSocketSession session : sessions) {
                if (!session.isOpen()) {
                    continue;
                }
                if (originSessionId != null && originSessionId.equals(session.getId())) {
                    // 排除掉发起者
                    continue;
                }
                try {
                    session.sendMessage(text);
                } catch (Exception e) {
                    log.warn("Failed to send ws message to session {}: {}", session.getId(), e.getMessage());
                }
            }
        } catch (Exception e) {
            log.error("Broadcast serialize/send failed: {}", e.getMessage(), e);
        }
    }

}
